﻿/*
Copyright 2011 Thomas W. Holtquist

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

using System;
using System.Collections.Generic;
using System.ComponentModel;
using System.Data;
using System.Drawing;
using System.IO;
using System.Linq;
using System.Reflection;
using System.Text;
using System.Windows.Forms;

namespace InstallMaker
{
    public partial class frmMain : Form
    {

        #region Variables

        private TreeNode nodeRootDir;       // Root Directory Node
        private TreeNode nodePrimaryPE;     // Node for the PE to install

        #endregion

        public frmMain()
        {
            InitializeComponent();

            // Setup Initial Directory Structure
            nodeRootDir = new TreeNode("\\Install Directory\\");
            trvStruct.Nodes.Add(nodeRootDir);
        }

        #region Object Methods

        /// <summary>
        /// Selects the primary PE for the project
        /// </summary>
        /// <param name="sender"></param>
        /// <param name="e"></param>
        private void btnBrowsePE_Click(object sender, EventArgs e)
        {
            // Prepare OpenFileDialog (OFD)
            OFD.Title = "Select Primary Application PE";
            OFD.Filter = "PE Files (*.PE)|*.PE";
            OFD.Multiselect = false;

            // Check for cancel
            if (OFD.ShowDialog() == System.Windows.Forms.DialogResult.Cancel) return;


            // Check for questionable files
            if (OFD.FileName.IndexOf("\\Debug\\") > 0)
            {
                if (MessageBox.Show("It appears you are using the Debug PE.\r\nAre you sure you wish to continue?", Application.ProductName, MessageBoxButtons.YesNo) == System.Windows.Forms.DialogResult.No) return;
            }

            if (OFD.FileName.IndexOf("\\be\\") > 0)
            {
                if (MessageBox.Show("It appears you are using the Big Endian PE.\r\nAre you sure you wish to continue?", Application.ProductName, MessageBoxButtons.YesNo) == System.Windows.Forms.DialogResult.No) return;
            }

            // Update text
            txtPE.Text = OFD.FileName;

            // Update structure
            if (nodePrimaryPE == null)
            {
                nodePrimaryPE = new TreeNode(Path.GetFileName(txtPE.Text).Replace(".pe",".pxe"));
                nodePrimaryPE.Tag = txtPE.Text;
                nodeRootDir.Nodes.Add(nodePrimaryPE);
                nodeRootDir.Expand();
            }
            else
            {
                nodePrimaryPE.Text = Path.GetFileName(txtPE.Text);
            }
        }

        /// <summary>
        /// Selects to directory to compile to
        /// </summary>
        /// <param name="sender"></param>
        /// <param name="e"></param>
        private void btnBrowseDIR_Click(object sender, EventArgs e)
        {
            FBD.Description = "Select a Build Directory";
            if (FBD.ShowDialog() == System.Windows.Forms.DialogResult.Cancel) return;
            txtOutput.Text = FBD.SelectedPath;
        }

        /// <summary>
        /// Add one or more files to the output structure
        /// </summary>
        /// <param name="sender"></param>
        /// <param name="e"></param>
        private void btnAddFile_Click(object sender, EventArgs e)
        {
            // Prepare OpenFileDialog (OFD)
            OFD.Title = "Select Primary Application PE";
            OFD.Filter = "All Files (*.*)|*.*|PE Files (*.PE)|*.PE|Images|*.bmp;*.jpg;*.jpeg;*.jpe;*.gif;*.wmf;*.png;*.ico";
            OFD.Multiselect = true;

            // Check for cancel
            if (OFD.ShowDialog() == System.Windows.Forms.DialogResult.Cancel) return;

            // Get the new directory's parent node
            TreeNode node = trvStruct.SelectedNode;
            if (node == null)
            {
                // Use the root directory
                node = nodeRootDir;
            }
            else
            {
                // Keep ascending until we reach a directory
                while (node.Text.Substring(0, 1) != "\\")
                    node = node.Parent;
            }

            // Add all files
            TreeNode newFile;
            for (int i = 0; i < OFD.FileNames.Length; i++)
            {
                newFile = new TreeNode(Path.GetFileName(OFD.FileNames[i]));
                newFile.Tag = OFD.FileNames[i];
                node.Nodes.Add(newFile);
            }
            node.Expand();

            btnRemAll.Enabled = true;
        }

        /// <summary>
        /// Adds a new directory to the output structure
        /// </summary>
        /// <param name="sender"></param>
        /// <param name="e"></param>
        private void btnAddDir_Click(object sender, EventArgs e)
        {
            InputBoxResult IBR = InputBox.Show("Please provide a name for your directory:", "Add Directory", string.Empty, null);
            if (!IBR.OK || IBR.Text == string.Empty) return;

            // Get the new directory's parent node
            TreeNode node = trvStruct.SelectedNode;
            if (node == null)
            {
                // Use the root directory
                node = nodeRootDir;
            }
            else
            {
                // Keep ascending until we reach a directory
                if (node.Text.Substring(0, 1) != "\\")
                {
                    while (node.Text.Substring(0, 1) != "\\")
                        node = node.Parent;
                }
            }

            // Add the node
            TreeNode newDir = new TreeNode("\\" + IBR.Text + "\\");
            node.Nodes.Add(newDir);
            node.Expand();

            // Select the new Directory
            trvStruct.SelectedNode = newDir;

            btnRemAll.Enabled = true;
        }

        /// <summary>
        /// Removes selected item from the Files/Folders treeview
        /// </summary>
        /// <param name="sender"></param>
        /// <param name="e"></param>
        private void btnRem_Click(object sender, EventArgs e)
        {
            trvStruct.SelectedNode.Remove();
            CheckRemAll();
        }

        /// <summary>
        /// Removes ALL items from the Files/Folders treeview (excepting root & primary pe)
        /// </summary>
        /// <param name="sender"></param>
        /// <param name="e"></param>
        private void btnRemAll_Click(object sender, EventArgs e)
        {
            TreeNode node;
            bool bFound;

            while (true)
            {
                bFound = false;
                for (int i = 0; i < nodeRootDir.Nodes.Count; i++)
                {
                    node = nodeRootDir.Nodes[i];
                    if (node != nodePrimaryPE)
                    {
                        node.Remove();
                        bFound = true;
                        break;
                    }
                }
                if (!bFound) break;
            }

            btnRemAll.Enabled = false;
        }

        /// <summary>
        /// Updates button enable based on selection
        /// </summary>
        /// <param name="sender"></param>
        /// <param name="e"></param>
        private void trvStruct_AfterSelect(object sender, TreeViewEventArgs e)
        {
            if (e.Node == null || e.Node == nodePrimaryPE || e.Node == nodeRootDir)
            {
                btnRem.Enabled = false;
            }
            else
            {
                btnRem.Enabled = true;
            }
        }

        /// <summary>
        /// Checks TreeView for ability to enable Remove All button
        /// </summary>
        private void CheckRemAll()
        {
            TreeNode node;

            for (int i = 0; i < nodeRootDir.Nodes.Count; i++)
            {
                node = nodeRootDir.Nodes[i];
                if (node != nodePrimaryPE)
                {
                    btnRemAll.Enabled = true;
                    return;
                }
            }

            btnRemAll.Enabled = false;
        }

        /// <summary>
        /// Adds a file association to the project
        /// </summary>
        /// <param name="sender"></param>
        /// <param name="e"></param>
        private void btnAddAssoc_Click(object sender, EventArgs e)
        {
            InputBoxResult IBR = InputBox.Show("What extension would you like to associate with this application?", "Add Association", string.Empty, null);
            if (!IBR.OK || IBR.Text == string.Empty) return;
            if (!IBR.Text.StartsWith(".")) IBR.Text = "." + IBR.Text;
            lstAssoc.Items.Add(IBR.Text);
            btnRemAllAssoc.Enabled = true;
        }

        /// <summary>
        /// Update build enabled based on PE, Output & Project texts
        /// </summary>
        /// <param name="sender"></param>
        /// <param name="e"></param>
        private void txtPE_TextChanged(object sender, EventArgs e)
        {
            mnuBuild.Enabled = (txtPE.Text != string.Empty && txtOutput.Text != string.Empty && txtProject.Text != string.Empty) ? true : false;
        }

        /// <summary>
        /// Update build enabled based on PE, Output & Project texts
        /// </summary>
        /// <param name="sender"></param>
        /// <param name="e"></param>
        private void txtOutput_TextChanged(object sender, EventArgs e)
        {
            mnuBuild.Enabled = (txtPE.Text != string.Empty && txtOutput.Text != string.Empty && txtProject.Text != string.Empty) ? true : false;
        }

        /// <summary>
        /// Updates remove association button based on selection
        /// </summary>
        /// <param name="sender"></param>
        /// <param name="e"></param>
        private void lstAssoc_SelectedIndexChanged(object sender, EventArgs e)
        {
            btnRemAssoc.Enabled = (lstAssoc.SelectedItem == null) ? false : true;
        }

        /// <summary>
        /// Removes the selected file association
        /// </summary>
        /// <param name="sender"></param>
        /// <param name="e"></param>
        private void btnRemAssoc_Click(object sender, EventArgs e)
        {
            lstAssoc.Items.Remove(lstAssoc.SelectedItem);
            btnRemAllAssoc.Enabled = (lstAssoc.Items.Count > 0) ? true : false;
        }

        /// <summary>
        /// Update build enabled based on PE, Output & Project texts
        /// </summary>
        /// <param name="sender"></param>
        /// <param name="e"></param>
        private void txtProject_TextChanged(object sender, EventArgs e)
        {
            mnuBuild.Enabled = (txtPE.Text != string.Empty && txtOutput.Text != string.Empty && txtProject.Text != string.Empty) ? true : false;
        }

        #endregion

        #region Build Methods

        /// <summary>
        /// Builds the install package
        /// </summary>
        /// <param name="sender"></param>
        /// <param name="e"></param>
        private void mnuBuild_Click(object sender, EventArgs e)
        {
            int i;
            string sOutput = ForceSlash(txtOutput.Text) + txtProject.Text.ToLower().Replace(" ", "") + ".pie";
            string sName = "";
            FileStream iFile;

            this.Enabled = false;
            this.Text = "InstallMaker - Building";

            // Delete the output if it exists
            if (File.Exists(sOutput))
                File.Delete(sOutput);

            // Create the FileStream
            iFile = new FileStream(sOutput, FileMode.CreateNew, FileAccess.Write);

            // Write the Header
            iFile.Write(new byte[] { (byte)'p', (byte)'i', (byte)'e', (byte)'_' }, 0, 4);   // Magic Number
            WriteName(txtProject.Text, iFile);                                              // Project Title

            // Write Root Files
            i = FilesInNode(nodeRootDir);                                                   // Files in Root Directory
            iFile.Write(new byte[] { (byte)i }, 0, 1);
            for (i = 0; i < nodeRootDir.Nodes.Count; i++)
            {
                if (nodeRootDir.Nodes[i].Text.Substring(0, 1) != "\\")
                {
                    sName =(string)nodeRootDir.Nodes[i].Tag; 
                    WriteFile(sName, sName.Replace(Path.GetFileName(sName),nodeRootDir.Nodes[i].Text), iFile);
                }
            }

            // Write Root Directories
            for (i = 0; i < nodeRootDir.Nodes.Count; i++)
            {
                if (nodeRootDir.Nodes[i].Text.Substring(0, 1) == "\\")
                    WriteDirectory("\\", nodeRootDir.Nodes[i], iFile);
            }

            // Terminate Directories
            iFile.Write(new byte[] { 0 }, 0, 1);                                            // Always terminate files with 0 byte

            // Write associations
            iFile.Write(new byte[] { (byte)lstAssoc.Items.Count }, 0, 1);
            for (i = 0; i < lstAssoc.Items.Count; i++)
                WriteName((string)lstAssoc.Items[i], iFile);
            
            iFile.Write(new byte[] { (byte)0 }, 0, 1);

            // Close the file
            iFile.Close();

            // End Build
            this.Text = "InstallMaker";
            MessageBox.Show("Your project has been built.");
            this.Enabled = true;

        }

        /// <summary>
        /// Forces a path to end in a slash
        /// </summary>
        /// <param name="path"></param>
        /// <returns></returns>
        private string ForceSlash(string path)
        {
            if (path.EndsWith("\\"))
                return path;
            return path + "\\";
        }

        /// <summary>
        /// Writes the name of a file/folder to the install file
        /// </summary>
        /// <param name="name"></param>
        /// <param name="fs"></param>
        private void WriteName(string name, FileStream fs)
        {
            if (name.Length > 255)
                throw new Exception("Directory or filename too long!\r\n255 character max, including parent directory name(s).");

            byte[] b = new byte[name.Length + 1];
            b[0] = (byte)name.Length;
            Array.Copy(Encoding.ASCII.GetBytes(name), 0, b, 1, b.Length - 1);
            fs.Write(b, 0, b.Length);
        }
        
        /// <summary>
        /// Gets the number of files in a supplied node
        /// </summary>
        /// <param name="topNode"></param>
        /// <returns></returns>
        private byte FilesInNode(TreeNode topNode)
        {
            int count = 0;
            int i;

            for (i = 0; i < topNode.Nodes.Count; i++)
            {
                if (topNode.Nodes[i].Text.Substring(0, 1) != "\\")
                    count++;
            }

            if (count > 255)
                throw new Exception("Too many files in directory '" + topNode.Text + "!\r\n255 files/directories allowed per directory.");

            return (byte)count;
        }

        /// <summary>
        /// Adds a directory to the install file
        /// </summary>
        /// <param name="root"></param>
        /// <param name="topNode"></param>
        /// <param name="fs"></param>
        private void WriteDirectory(string root, TreeNode topNode, FileStream fs)
        {
            int i;

            this.Text = "InstallMaker - Building (" + (root + topNode.Text.Substring(1, topNode.Text.Length - 2)) + "\\)";

            // Write the directory name
            WriteName(root + topNode.Text.Substring(1, topNode.Text.Length - 2), fs);

            // Write Files
            i = FilesInNode(topNode);
            fs.Write(new byte[] { (byte)i }, 0, 1);
            for (i = 0; i < topNode.Nodes.Count; i++)
            {
                if (topNode.Nodes[i].Text.Substring(0, 1) != "\\")
                    WriteFile((string)topNode.Nodes[i].Tag, fs);
            }

            // Write Root Directories
            for (i = 0; i < topNode.Nodes.Count; i++)
            {
                if (topNode.Nodes[i].Text.Substring(0, 1) == "\\")
                    WriteDirectory(root + topNode.Text, topNode.Nodes[i], fs);
            }
        }

         /// <summary>
        /// Writes a file to the install file
        /// </summary>
        /// <param name="file"></param>
        /// <param name="outputName"></param>
        /// <param name="fs"></param>
        private void WriteFile(string file, FileStream fs)
        {
            WriteFile(file, file, fs);
        }

        /// <summary>
        /// Writes a file to the install file with a specified name
        /// </summary>
        /// <param name="file"></param>
        /// <param name="outputName"></param>
        /// <param name="fs"></param>
        private void WriteFile(string file, string outputName, FileStream fs)
        {
            byte[] b;
            long size = new System.IO.FileInfo(file).Length;

            this.Text = "InstallMaker - Building (" + file + ")";
            WriteName(Path.GetFileName(outputName), fs);      // Write the File Name

            // Write File Size
            b = new byte[4];
            b[0] = (byte)size;
            b[1] = (byte)(size >> 8);
            b[2] = (byte)(size >> 16);
            b[3] = (byte)(size >> 24);
            fs.Write(b, 0, 4);

            // Write the File
            fs.Write(File.ReadAllBytes(file), 0, (int)size);
        }

        #endregion

        #region Test Methods

        /// <summary>
        /// Test an install file
        /// </summary>
        /// <param name="sender"></param>
        /// <param name="e"></param>
        private void mnuTest_Click(object sender, EventArgs e)
        {
            byte[] b, bOut;
            long size;
            string sInstall, sPath, sName;
            string sProject = "";
            long iIndex = 4;
            int iLen, i, j, iCount;
            FileStream iFile;

            // Prepare OpenFileDialog (OFD)
            OFD.Title = "Select an installer to test";
            OFD.Filter = "InstallMaker Files|*.pie";
            OFD.Multiselect = false;

            // Check for cancel
            if (OFD.ShowDialog() == System.Windows.Forms.DialogResult.Cancel) return;

            // Read File
            size = new System.IO.FileInfo(OFD.FileName).Length;
            b = new byte[size];
            b = File.ReadAllBytes(OFD.FileName);

            if (b.Length < 5)
            {
                MessageBox.Show("Invalid file length!");
                return;
            }

            // Check header
            if (b[0] != 112 || b[1] != 105 || b[2] != 101 || b[3] != 95)
            {
                MessageBox.Show("Invalid header!");
                return;
            }

            this.Text = "InstallMaker - Testing";

            // Read project name
            iLen = b[iIndex++];
            for (i = 0; i < iLen; i++)
                sProject = string.Concat(sProject, (char)b[iIndex++]);

            // Select Install Directory
            FBD.Description = "Select an Install Directory for " + sProject;
            if (FBD.ShowDialog() == System.Windows.Forms.DialogResult.Cancel) return;
            sInstall = FBD.SelectedPath;
            if (!sInstall.EndsWith("\\")) sInstall += "\\";

            // Install Root Files
            try
            {
                iCount = b[iIndex++];   // Check the number of items in the directory
                for (i = 0; i < iCount; i++)
                {
                    // Read file name
                    sName = "";
                    iLen = b[iIndex++];
                    for (j = 0; j < iLen; j++)
                        sName = string.Concat(sName, (char)b[iIndex++]);

                    this.Text = "InstallMaker - Testing [" + sName + "]";

                    // Get File Size
                    iLen = ((b[iIndex + 3] & 0xFF) << 24) + ((b[iIndex + 2] & 0xFF) << 16) + ((b[iIndex + 1] & 0xFF) << 8) + (b[iIndex] & 0xFF);
                    iIndex += 4;

                    // Delete existant file
                    if (File.Exists(sInstall + sName))
                        File.Delete(sInstall + sName);

                    // Write File
                    bOut = new byte[iLen];
                    Array.Copy(b, iIndex, bOut, 0, iLen);
                    iFile = new FileStream(sInstall + sName, FileMode.CreateNew, FileAccess.Write);
                    iFile.Write(bOut, 0, bOut.Length);
                    iFile.Close();
                    iIndex += iLen;
                }
            }
            catch (Exception)
            {
                MessageBox.Show("Install failed!");
                this.Text = "InstallMaker";
            }

            // Install Remaining Folders & Files
            try
            {
                while (true)
                {
                    // Read Directory Name or EOF
                    sName = "";
                    iLen = b[iIndex++];
                    if (iLen == 0) break;
                    for (i = 0; i < iLen; i++)
                        sName = string.Concat(sName, (char)b[iIndex++]);
                    if (!sName.EndsWith("\\")) sName += "\\";

                    // Create Directory
                    sPath = sInstall + sName;
                    if (!Directory.Exists(sPath)) Directory.CreateDirectory(sPath);
                    this.Text = "InstallMaker - Testing [" + sPath + "]";

                    // Write files in Directory
                    iCount = b[iIndex++];
                    for (i = 0; i < iCount; i++)
                    {
                        // Read file name
                        sName = "";
                        iLen = b[iIndex++];
                        for (j = 0; j < iLen; j++)
                            sName = string.Concat(sName, (char)b[iIndex++]);
                        this.Text = "InstallMaker - Testing [" + sName + "]";

                        // Get File Size
                        iLen = ((b[iIndex + 3] & 0xFF) << 24) + ((b[iIndex + 2] & 0xFF) << 16) + ((b[iIndex + 1] & 0xFF) << 8) + (b[iIndex] & 0xFF);
                        iIndex += 4;

                        // Delete existant file
                        if (File.Exists(sPath + sName))
                            File.Delete(sPath + sName);

                        // Write File
                        bOut = new byte[iLen];
                        Array.Copy(b, iIndex, bOut, 0, iLen);
                        iFile = new FileStream(sPath + sName, FileMode.CreateNew, FileAccess.Write);
                        iFile.Write(bOut, 0, bOut.Length);
                        iFile.Close();
                        iIndex += iLen;
                    }

                }
            }
            catch (Exception)
            {
                MessageBox.Show("Install failed!");
                this.Text = "InstallMaker";
            }


            MessageBox.Show("Install completed OK");
            this.Text = "InstallMaker";
        }

        #endregion

    }
}
